use arrow::array::{Array, ArrayRef, Float64Array};
use anyhow::{Context, Result};

pub fn calculate_ma(values: &Float64Array, window: usize) -> Result<Float64Array> {
    let mut ma = Vec::with_capacity(values.len());
    
    // 填充前window-1个值
    for _ in 0..window-1 {
        ma.push(0.0);
    }

    for i in window-1..values.len() {
        let sum: f64 = values
            .values()
            .slice(i - window + 1, window)
            .iter()
            .sum();
        ma.push(sum / window as f64);
    }

    Ok(Float64Array::from(ma))
}

pub fn calculate_ema(values: &Float64Array, window: usize) -> Result<Float64Array> {
    let mut ema = Vec::with_capacity(values.len());
    let alpha = 2.0 / (window as f64 + 1.0);
    
    // 第一个值使用简单平均
    let first_ma = values
        .values()
        .slice(0, window.min(values.len()))
        .iter()
        .sum::<f64>() / window.min(values.len()) as f64;
    ema.push(first_ma);

    // 计算其余值
    for i in 1..values.len() {
        let current = values.value(i);
        let previous_ema = ema[i - 1];
        let new_ema = current * alpha + previous_ema * (1.0 - alpha);
        ema.push(new_ema);
    }

    Ok(Float64Array::from(ema))
}

pub fn calculate_rsi(values: &Float64Array, window: usize) -> Result<Float64Array> {
    let mut rsi = Vec::with_capacity(values.len());
    let mut gains = Vec::new();
    let mut losses = Vec::new();

    // 计算价格变化
    for i in 1..values.len() {
        let change = values.value(i) - values.value(i - 1);
        gains.push(change.max(0.0));
        losses.push((-change).max(0.0));
    }

    // 填充第一个值
    rsi.push(50.0);

    // 计算初始平均收益和损失
    let mut avg_gain = gains.iter().take(window).sum::<f64>() / window as f64;
    let mut avg_loss = losses.iter().take(window).sum::<f64>() / window as f64;

    // 计算RSI
    for i in window..values.len() - 1 {
        // 更新平均收益和损失
        avg_gain = (avg_gain * (window - 1) as f64 + gains[i]) / window as f64;
        avg_loss = (avg_loss * (window - 1) as f64 + losses[i]) / window as f64;

        // 计算相对强度
        let rs = if avg_loss == 0.0 {
            100.0
        } else {
            avg_gain / avg_loss
        };

        // 计算RSI
        let rsi_value = 100.0 - (100.0 / (1.0 + rs));
        rsi.push(rsi_value);
    }

    Ok(Float64Array::from(rsi))
}

pub fn calculate_bollinger_bands(
    values: &Float64Array,
    window: usize,
    num_std: f64,
) -> Result<(Float64Array, Float64Array, Float64Array)> {
    let middle = calculate_ma(values, window)?;
    let mut upper = Vec::with_capacity(values.len());
    let mut lower = Vec::with_capacity(values.len());

    // 填充前window-1个值
    for i in 0..window-1 {
        upper.push(middle.value(i));
        lower.push(middle.value(i));
    }

    for i in window-1..values.len() {
        let window_slice: Vec<f64> = values
            .values()
            .slice(i - window + 1, window)
            .iter()
            .copied()
            .collect();

        let mean = middle.value(i);
        let variance = window_slice
            .iter()
            .map(|&x| (x - mean).powi(2))
            .sum::<f64>() / (window - 1) as f64;
        let std_dev = variance.sqrt();

        upper.push(mean + num_std * std_dev);
        lower.push(mean - num_std * std_dev);
    }

    Ok((
        Float64Array::from(upper),
        middle,
        Float64Array::from(lower),
    ))
}

pub fn calculate_atr(
    high: &Float64Array,
    low: &Float64Array,
    close: &Float64Array,
    window: usize,
) -> Result<Float64Array> {
    let mut tr = Vec::with_capacity(high.len());
    let mut atr = Vec::with_capacity(high.len());

    // 计算真实范围（TR）
    tr.push(high.value(0) - low.value(0)); // 第一个TR只使用当日高低差

    for i in 1..high.len() {
        let h_l = high.value(i) - low.value(i);
        let h_pc = (high.value(i) - close.value(i - 1)).abs();
        let l_pc = (low.value(i) - close.value(i - 1)).abs();
        tr.push(h_l.max(h_pc).max(l_pc));
    }

    // 计算ATR
    // 第一个ATR使用简单平均
    let first_atr = tr.iter().take(window).sum::<f64>() / window as f64;
    atr.push(first_atr);

    // 计算剩余的ATR值
    for i in window..tr.len() {
        let previous_atr = atr[i - 1];
        let current_tr = tr[i];
        let new_atr = (previous_atr * (window - 1) as f64 + current_tr) / window as f64;
        atr.push(new_atr);
    }

    Ok(Float64Array::from(atr))
}